import numpy as np
import openai
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langfuse.decorators import observe, langfuse_context
# 加载 .env 到环境变量
from dotenv import load_dotenv, find_dotenv
from langfuse.decorators import observe, langfuse_context

_ = load_dotenv(find_dotenv())
''' 
 这个实际应用就是根据已有的内容，然后根据用户提出的问题判断是否在改内容之内，如果在就返回true，否则返回false
 每次调用大模型之前，要先判断这个问题是否已经有人提出了，如果有人提出了，则直接返回true，否则加入问题列表

'''

# 需要回答问题的prompt
need_answer = PromptTemplate.from_template(
    """
    *********
    你是AIGC课程的助教，你的工作是从学员的课堂交流中选择出需要老师回答的问题，加以整理以交给老师回答。

    课程内容:
    {outlines}
    *********
    学员输入:
    {user_input}
    *********
    如果这是一个需要老师答疑的问题，回复Y，否则回复N。
    只回复Y或N，不要回复其他内容。
    """
)

# 检测提出的问题是否已经有答案了
check_duplicated = PromptTemplate.from_template(
    """
*********
已有提问列表:
[
{question_list}
]
*********
新提问:
{user_input}
*********
已有提问列表是否有和新提问类似的问题? 回复Y或N, Y表示有，N表示没有。
只回复Y或N，不要回复其他内容。
    """
)
# 课程内容
outlines = """
LangChain
模型 I/O 封装
模型的封装
模型的输入输出
PromptTemplate
OutputParser
数据连接封装
文档加载器：Document Loaders
文档处理器
内置RAG：RetrievalQA
记忆封装：Memory
链架构：Chain/LCEL
大模型时代的软件架构：Agent
ReAct
SelfAskWithSearch
LangServe
LangChain.js
"""
parser = StrOutputParser()
# 创建 chain
model = ChatOpenAI(temperature=0, model_kwargs={"seed": 42})
need_answer_chain = (
        need_answer
        | model
        | parser
)
# 是否需要回答 chain
is_duplicated_chain = (
        check_duplicated
        | model
        | parser
)
# 对prompt进行缓存
cache = {}


@observe()
def get_embeddings(text):
    '''封装 OpenAI 的 Embedding 模型接口'''
    if text in cache:
        # 如果已经在缓存中，不再重复调用（节省时间、费用）
        return cache[text]
    # 否则就去embeddings大模型查询问题
    data = openai.embeddings.create(
        input=text,
        model="text-embedding-3-small",
        dimensions=256
    ).data
    cache[text] = data[0].embedding
    return data[0].embedding


@observe()
def cos_sim(v, m):
    '''计算cosine相似度'''
    score = np.dot(m, v) / (np.linalg.norm(m, axis=1) * np.linalg.norm(v))
    return score.tolist()


@observe()
# 相识度超过threshold就代表是重复问题
def check_duplicated(query, existing, threshold=0.825):
    '''通过cosine相似度阈值判断是否重复'''
    query_vec = np.array(get_embeddings(query))
    mat = np.array([item[1] for item in existing])
    # 计算相识度
    cos = cos_sim(query_vec, mat)
    # 大于等于阈值就代表问题重复
    return max(cos) >= threshold


@observe()
def need_answer(question, outlints):
    '''判断是否需要回答'''
    langfuse_handler = langfuse_context.get_current_langchain_handler()
    return need_answer_chain.invoke(
        {"user_input": question, "outlines": outlints},
        config={"callbacks": [langfuse_handler]}
    ) == 'Y'


# 假设 已有问题
question_list = [
    ("LangChain可以商用吗", get_embeddings("LangChain可以商用吗")),
    ("LangChain开源吗", get_embeddings("LangChain开源吗")),
]


@observe()
def verify_question(
        question: str,
        outlines: str,
        question_list: list,
        user_id,
) -> bool:
    langfuse_context.update_current_trace(
        name="AGIClassAssistant2",
        user_id=user_id,
    )
    # 判断是否需要回答
    if need_answer(question, outlines):
        # 判断问题是否重复
        if not check_duplicated(question, question_list):
            # 代表问题重复
            vec = cache[question]
            question_list.append((question, vec))
            return True
    return False


ret = verify_question(
    "LangChain支持Java吗",
    # "LangChain有商用许可吗",
    outlines,
    question_list,
    user_id="wzr"
)
print(ret)
